from langchain_openai import ChatOpenAI
from langchain.agents import Tool, initialize_agent, AgentType
from server.config import OPENAI_API_KEY, OPENAI_BASE_URL

llm = ChatOpenAI(
    temperature=0.3,
    base_url=OPENAI_BASE_URL,
    model="gpt-3.5-turbo",
    openai_api_key=OPENAI_API_KEY,
)


# Tool 1：代码补全
def code_completion(input_code: str) -> str:
    prompt = f"请补全下面的算法竞赛代码：\n{input_code}\n补全后的完整代码是："
    return llm.invoke(prompt).content


# Tool 2：题目思路分析
def problem_solver(problem_description: str) -> str:
    prompt = (
        f"下面是一道算法竞赛题，请帮我分析解题思路：\n{problem_description}\n思路分析："
    )
    return llm.invoke(prompt).content


# 构造工具
tools = [
    Tool.from_function(
        func=code_completion,
        name="CodeCompleter",
        description="根据输入的不完整代码补全完整代码，适用于算法竞赛。",
    ),
    Tool.from_function(
        func=problem_solver,
        name="ProblemSolver",
        description="根据输入的题目描述输出解题思路，适用于算法竞赛题分析。",
    ),
]

# 创建 Agent
agent = initialize_agent(
    tools=tools, llm=llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True
)
